import numpy as np
from sklearn.datasets import load_digits
from sklearn.preprocessing import LabelBinarizer
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
import matplotlib.pyplot as plt

# 载入数据
digits = load_digits()
# 显示图片
for i in range(min(digits.images.shape[0], 2)):
    plt.imshow(digits.images[i], cmap='gray')
    plt.show()

# 数据
X = digits.data
# 标签
y = digits.target

# 定义一个神经网络，结构，64-100-
# 定义输入层到隐藏层之间的权值矩阵
V = np.random.random((64, 100)) * 2 - 1
# 定义隐藏层到输出层之间的权值矩阵
W = np.random.random((100, 10)) * 2 - 1

# 数据切分
# 1/4为测试集，3/4为训练集
X_train, X_test, y_train, y_test = train_test_split(X, y)

# 标签二值化
# 0 -> 1000000000
# 3 -> 0003000000
# 9 -> 0000000001
labels_train = LabelBinarizer().fit_transform(y_train)


# 激活函数
def sigmoid(x):
    return 1 / (1 + np.exp(-x))


# 激活函数的导数
def dsigmoid(x):
    return x * (1 - x)


#  训练模型
def train(X, y, steps=10000, lr=0.011):
    global V, W
    for n in range(steps + 1):
        # 随机选取一个数据
        i = np.random.randint(X.shape[0])
        # 获取一个数据
        x = X[i]
        x = np.atleast_2d(x)
        # BP算法公式
        # 计算隐藏层的输出
        L1 = sigmoid(np.dot(x, V))
        # 计算输出层的输出
        L2 = sigmoid(np.dot(L1, W))
        # 计算L2_delta,L1_delta
        L2_delta = (y[i] - L2) * dsigmoid(L2)
        L1_delta = L2_delta.dot(W.T) * dsigmoid(L1)
        # 更新权值
        W += lr * L1.T.dot(L2_delta)
        V += lr * x.T.dot(L1_delta)

        # 每训练1000次预测一次准确率
        if n % 1000 == 0:
            output = predict(X_test)
            predictions = np.argmax(output, axis=1)
            acc = np.mean(np.equal(predictions, y_test))
            dW = L1.T.dot(L2_delta)
            dV = x.T.dot(L1_delta)
            gradient = np.sum([np.sqrt(np.sum(np.square(j))) for j in [dW, dV]])
            print('steps', n, 'accuracy', acc, 'gradient', gradient)
            # print(classification_report(predictions,y_test))


def predict(x):
    # 计算隐藏层的输出
    L1 = sigmoid(np.dot(x, V))
    # 计算输出层的输出
    L2 = sigmoid(np.dot(L1, W))
    return L2


# 开始训练(设置迭代次数)
train(X_train, labels_train, 100000, lr=0.11)
train(X_train, labels_train, 100000, lr=0.011)
# 训练后结果对比
output = predict(X_test)
predictions = np.argmax(output, axis=1)
acc = np.mean(np.equal(predictions, y_test))
print('accuracy', acc)
print(classification_report(predictions, y_test, digits=4))

